Si Yuan Lee | SWE | Building ML Infra for Recommendation

Triton Beginner Guide

Triton is a language and compiler for writing high-performance GPU kernels with less boilerplate than CUDA while still exposing the parts of the hardware that matter for performance. To use it well, it helps to understand how GPU memory, execution, and transformer workloads fit together.

GPU Memory Hierarchy

Fastest to slowest:

GPU SRAM: 19 TB/s, 20 MB
GPU HBM: 1.5 TB/s, 40 GB
CPU DRAM: slowest in this hierarchy

The practical lesson is simple: performance usually improves when data stays as close as possible to the compute units. A good kernel tries to reuse data from on-chip memory instead of repeatedly going back to HBM or CPU memory.

Memory Flow Diagram

CPU DRAM
   |
   v
GPU HBM
   |
   v
SM-local SRAM / shared memory / registers
   |
   v
Arithmetic units execute the kernel

Core Concepts

A streaming multiprocessor (SM) is the unit that schedules and executes groups of threads.
Each SM has access to fast on-chip storage such as registers and shared memory or SRAM.
A useful mental model is that a GPU core contributes floating-point work, although actual throughput depends on instruction mix, tensor cores, memory access, and occupancy.
A warp is a group of threads that execute in lockstep.
- 32 threads on NVIDIA GPUs such as H100 and A100
- 64 threads on many AMD architectures, often called a wavefront
A warp executes the same instruction at the same time across its lanes, which is why branch divergence can hurt performance.

Why This Matters for Triton

Triton kernels are typically written around tiles or blocks of work.
Those tiles are chosen so the kernel can load data once, reuse it many times, and reduce slow memory traffic.
This is especially important for matrix multiplication, attention, layer normalization, and other bandwidth-sensitive operations.

Decoder-Only Transformer Architecture

Decoder-only transformers are the foundation behind many modern language models. They are optimized for autoregressive generation, where the model predicts the next token from all previous tokens.

Why Decoder-Only

Trains on next-token prediction, which fits a decoder-only architecture with a causal mask.
There is no need for paired data such as $X \to Y$ .
KV cache improves efficiency for K and V because these entries are reused for each query during inference for MHA(Q, K, V).
Without cache: $O(n^2)$
With cache: approximately $O(n)$ work per decoding step for the newly added token

During training, the model sees a full token sequence and learns to predict the next token at each position. During inference, generation happens one token at a time, so caching old keys and values prevents recomputing the entire attention history for every new token.

High-Level Decoder Block

Input x
  |
  v
LayerNorm
  |
  v
Masked Multi-Head Self-Attention
  |
  v
Residual Add
  |
  v
LayerNorm
  |
  v
Feed-Forward Network
  |
  v
Residual Add
  |
  v
Output x

Block Structure

x
|
+---------------------------------------------+
|                                             |
v                                             |
+------------------+                          |
|    LayerNorm     |                          |
+------------------+                          |
      |                                    |
      v                                    |
+-------------------------------+             |
| Masked Multi-Head Attention   |             |
+-------------------------------+             |
      |                                    |
      v                                    |
+-------------------------------+             |
| Add residual connection from x| <-----------+
+-------------------------------+
      |
      v
     x'
      |
      +---------------------------------------------+
      |                                             |
      v                                             |
+------------------+                                   |
|    LayerNorm     |                                   |
+------------------+                                   |
      |                                             |
      v                                             |
+-------------------------------+                      |
|     Feed-Forward Network      |                      |
+-------------------------------+                      |
      |                                             |
      v                                             |
+--------------------------------+                     |
| Add residual connection from x'| <-------------------+
+--------------------------------+
      |
      v
   output

This residual structure is important because it stabilizes optimization and allows very deep stacks of layers. In practice, most of the expensive operations are concentrated in attention and the feed-forward network.

Attention Diagram

Input embeddings
    |
    +--> Q projection
    +--> K projection
    +--> V projection
             |
             v
  Attention(Q, K, V) with causal mask
             |
             v
       Concatenate heads
             |
             v
        Output projection

Formulas

\operatorname{Attention}(Q, K, V) = \operatorname{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + \text{mask}\right)V

\operatorname{FFN}(x) = W_2 \, \sigma(W_1x + b_1) + b_2

x' = x + \operatorname{Attention}(\operatorname{LN}(x))

\operatorname{Block}(x) = x' + \operatorname{FFN}(\operatorname{LN}(x'))

The exact activation sigma depends on the model family. Common choices include ReLU, GELU, and SwiGLU variants. For modern LLMs, the feed-forward block is often one of the largest consumers of compute and memory bandwidth.

Why KV Cache Changes Inference Cost

Without KV cache, every generated token would recompute attention over the entire prefix from scratch. With KV cache:

old K and V are stored once
only the new token produces a new Q, K, and V
attention still reads previous tokens, but it avoids recomputing their projections

This is why the per-step cost becomes much closer to linear in the current context length rather than rebuilding the entire sequence state each time.

Where Triton Optimization Fits

Common optimization techniques include flash attention, fused kernels, and tiling.
Custom Triton kernels can still provide efficiency gains, especially for short sequence lengths.

In practice, Triton is most useful when you care about one of the following:

Kernel fusion: combine multiple small operations to reduce memory traffic.
Tiling: load a block of data into fast memory, do more work on it, then write results back.
Better memory access patterns: improve coalescing and reduce wasted bandwidth.
Shape-specific optimization: specialize for sizes that generic libraries do not handle optimally.
Short-sequence workloads: reduce overhead where highly general kernels may be underutilized.

Optimization Map

Transformer workload
    |
    +--> Attention ---------> Flash attention / tiling / SRAM reuse
    +--> FFN ---------------> Matmul blocking / fusion / vectorized loads
    +--> LayerNorm ---------> Reduction patterns / fusion
    +--> Elementwise ops ---> Fuse into neighboring kernels

The core idea is that Triton does not replace the transformer architecture. It targets the implementation layer, where execution order, tile size, memory movement, and fusion strategy determine whether the hardware is fully utilized.

Final Intuition

If you remember only one thing, it should be this: modern GPU performance is often limited less by raw arithmetic and more by how efficiently data moves through memory. Triton is valuable because it gives you direct control over that data movement without forcing you to write extremely low-level GPU code.